Implement Agent Trajectory as a formal data structure.#969
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughChangesThe PR adds typed trajectory storage with CSV and JSONL persistence, integrates it into environment and DSE execution, and updates reporting to load JSONL trajectories with CSV fallback. Tests cover trajectory behavior, caching, environment parameters, persistence, and reporting. Trajectory integration
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/cloudai/configurator/trajectory.py`:
- Around line 32-55: Make the identity-bearing data in
EnvParamsSample.env_params and TrialResult.action structurally immutable, not
only field-reassignment immutable, so values returned through
TrajectoryEntry.get() cannot alter cache matching. Deep-freeze nested mappings
and sequences when constructing or storing these components, or maintain a
private immutable canonical key used by find(); ensure persisted records and
future find() calls use that stable identity.
- Around line 104-118: Update TrajectoryWriter.append to inspect an existing
trajectory.csv header before accepting the first record schema: compare any
non-empty header with self._fields and raise ValueError on mismatch, while
allowing a pre-created empty file to be initialized with the header. Ensure
subsequent appends continue enforcing the established field order and do not
write duplicate headers.
In `@src/cloudai/systems/slurm/single_sbatch_runner.py`:
- Around line 208-225: Update unroll_dse in SingleSbatchRunner to support
declared tr.test.env_params by sampling an EnvParamsSample for each trial,
applying those same values to the execution test run and gym trajectory append.
Ensure the sampled environment parameters are passed wherever CloudAIGymEnv
requires them, and add an integration test covering DSE actions with environment
parameters.
In `@tests/test_cloudaigym.py`:
- Line 717: Suppress Ruff’s S311 finding for the deterministic Random usage in
the expected_sample setup, using the repository’s standard inline suppression
syntax and keeping the test’s deterministic behavior unchanged.
In `@tests/test_reporter.py`:
- Around line 50-59: Update test_load_trajectory_dataframe_prefers_json_lines to
also create a conflicting trajectory.csv in tmp_path with different data before
calling load_trajectory_dataframe. Keep the existing JSONL assertions and verify
the returned path and records still come from trajectory.jsonl, proving JSONL
takes precedence over CSV.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 469cd551-c61f-434b-afff-96cfbef6bf41
📒 Files selected for processing (13)
src/cloudai/configurator/__init__.pysrc/cloudai/configurator/cloudai_gym.pysrc/cloudai/configurator/env_params.pysrc/cloudai/configurator/trajectory.pysrc/cloudai/models/workload.pysrc/cloudai/report_generator/dse_report.pysrc/cloudai/reporter.pysrc/cloudai/systems/slurm/single_sbatch_runner.pytests/test_cloudaigym.pytests/test_env_params.pytests/test_gymnasium_adapter_contract.pytests/test_reporter.pytests/test_trajectory.py
💤 Files with no reviewable changes (2)
- src/cloudai/configurator/env_params.py
- tests/test_env_params.py
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/cloudai/configurator/trajectory.py`:
- Around line 258-264: The trajectory entry storage used by find() must snapshot
identity-contributing component state instead of comparing live mutable
instances. Update the append/store path and _identity_for handling so every
component with contributes_to_identity=True is frozen or represented by a
private immutable canonical key, while preserving existing identity matching for
persisted entries.
In `@src/cloudai/systems/slurm/single_sbatch_runner.py`:
- Around line 217-233: Apply the same constraint_check used by unroll_dse() in
the combination loop before creating the trial observation, computing reward, or
appending trajectory data; skip rejected combinations so only submitted runs are
recorded. Add a regression test covering a constraint-rejected DSE combination
and verifying it is absent from trajectory output and reports.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 3c0c2a29-883b-4460-86f3-6aba6ecf9382
📒 Files selected for processing (6)
src/cloudai/configurator/trajectory.pysrc/cloudai/systems/slurm/single_sbatch_runner.pytests/test_cloudaigym.pytests/test_reporter.pytests/test_single_sbatch_runner.pytests/test_trajectory.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/cloudai/configurator/trajectory.py`:
- Around line 352-356: Update _validate_components to reject any component
dataclass with contributes_to_identity=True unless its dataclass definition is
frozen. Preserve existing validation for non-dataclass components and duplicate
types, and raise a clear TypeError identifying the offending component.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: b339fc85-47e7-4b63-a8c0-19754eb9876c
📒 Files selected for processing (3)
src/cloudai/configurator/trajectory.pysrc/cloudai/systems/slurm/single_sbatch_runner.pytests/test_trajectory.py
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/cloudai/configurator/trajectory.py (1)
43-53: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winSnapshot the observation payload before storing it.
TrialResultfreezesactionbut leavesobservationcaller-owned. A mutable list can be changed after_store_entry()persists the row, causing cache hits to return observations that differ from the persisted trajectory. Freeze or deep-copy the observation as well.Proposed fix
def __post_init__(self) -> None: """Own an immutable snapshot of the action.""" object.__setattr__(self, "action", _freeze(self.action)) + object.__setattr__(self, "observation", _freeze(self.observation))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cloudai/configurator/trajectory.py` around lines 43 - 53, Update TrialResult.__post_init__ to snapshot observation before storing it, using the existing _freeze helper or an equivalent deep-copy mechanism. Preserve the observation’s sequence semantics while ensuring later caller mutations cannot alter the stored trajectory or cached results.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/cloudai/configurator/base_agent.py`:
- Line 52: Update the Slurm DSE environment construction in
single_sbatch_runner.py, where CloudAIGymEnv is instantiated, to pass through
the configured trajectory_file_type alongside rewards. Preserve the selected
value, including "jsonl", instead of allowing the environment default of "csv"
to apply.
---
Outside diff comments:
In `@src/cloudai/configurator/trajectory.py`:
- Around line 43-53: Update TrialResult.__post_init__ to snapshot observation
before storing it, using the existing _freeze helper or an equivalent deep-copy
mechanism. Preserve the observation’s sequence semantics while ensuring later
caller mutations cannot alter the stored trajectory or cached results.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 365ca573-91b2-4412-be87-68fada14bbbd
📒 Files selected for processing (3)
src/cloudai/cli/handlers.pysrc/cloudai/configurator/base_agent.pysrc/cloudai/configurator/trajectory.py
Summary
Introduces a typed Trajectory data structure with extensible dataclass components, cache identity support, and CSV/JSONL persistence. Environment parameters are now stored as trajectory components, and CloudAI defaults to CSV output.
Features
The new Trajectory class manages in-memory trajectory entries and persists them to a file: trajectory.[csv/jsonl].
Initialize it with the iteration directory and any optional component types:
Append a completed trial using the fields required by the configured components:
Entries can be accessed as an ordered sequence or searched by trial identity:
Trajectory file type output can be selected using the new agent_config flag.
Test Plan
Tested locally with Python 3.14 using:
Also tested with CloudAI internal agents.
Notes